import Router from 'koa-router';
import getOpenAIInstance from './lib/openai.js';

const router = new Router();

router.get('/api/gpt/chat', async (ctx, next) => {
	ctx.status = 200;
	ctx.set('Content-Type', 'text/event-stream'); // 'text/event-stream' 标识 SSE 即 Server-Sent Events

	let gptStream;

	// req on close 这段代码应该写到前面，不要写到最后，写到最后的话，只能是等待 res.write 输出完才能执行，当然中途监听不到了。

	ctx.req.on('close', () => {
		console.log('req close...');

		// 取消请求
		console.log('try abort...');

		if (!ctx.gptStreamDone) {
			console.log('abort request...');
			// 通过查阅 openai 文档(https://www.npmjs.com/package/openai) 可以使用 gptStream.controller.abort() 来中止请求。所以核心代码可以这样写：
			gptStream.controller.abort();
			console.log('abort ok~ ');
		}
	});

	const query = ctx.query || {};

	// 简单的密钥
	const authToken = query['x-auth-token'] || '';
	if (!authToken.trim() || authToken !== process.env.AUTH_TOKEN) {
		const errMsg = 'invalid token';
		console.log('error: ', errMsg);
		ctx.res.write(`data: [ERROR]${errMsg}\n\n`); // 格式必须是 `data: xxx\n\n` ！！！
		return;
	}

	// option
	const optionStr = query['option'] || '{}';
	const decodeOptionStr = decodeURIComponent(optionStr);
	const option = JSON.parse(decodeOptionStr);

	if (!option.messages) {
		const errMsg = 'invalid option: messages required';
		console.log('error: ', errMsg);
		ctx.res.write(`data: [ERROR]${errMsg}\n\n`);
		return;
	}

	// get openai instance
	const instance = getOpenAIInstance();
	if (instance == null) {
		const errMsg = 'openai instance is null';
		console.log('error: ', errMsg);
		ctx.res.write(`data: [ERROR]${errMsg}\n\n`);
		return;
	}

	const formatKey = instance.key.slice(0, 20) + '***';
	console.log('cur openai key: ', formatKey);

	try {
		// request GPT API
		gptStream = await instance.openai.chat.completions.create({
			model: instance.model,
			// messages: [{ role: 'user', content: 'xxx' }],
			max_tokens: 600,
			stream: true, // stream
			stream_options: { include_usage: true },
			...option
		});
	} catch (err) {
		const errMsg = err.message || 'request openai API error';
		console.log('error: ', errMsg);
		ctx.res.write(`data: [ERROR]${errMsg}\n\n`);

		// 记录 error 并发送邮件
		instance.isError = true;
		sendEmail({
			subject: `OpenAI API request error, key ${formatKey}`,
			text: errMsg
		});

		return;
	}

	if (gptStream == null) {
		const errMsg = 'gptStream is not defined';
		console.log('error: ', errMsg);
		ctx.res.write(`data: [ERROR]${errMsg}\n\n`);
		return;
	}

	for await (const chunk of gptStream) {
		const { choices = [], usage } = chunk;

		if (choices.length === 0 || usage != null) {
			// console.log('usage ', usage) // 格式如 {"prompt_tokens":81,"completion_tokens":23,"total_tokens":104}
			ctx.res.write(`data: ${JSON.stringify({ usage })}\n\n`);

			ctx.gptStreamDone = true;
			ctx.res.write(`data: [DONE]\n\n`);
		}

		if (choices.length > 0) {
			const content = choices[0].delta.content;
			console.log('content: ', content);

			if (content != null) {
				const data = { c: content };
				ctx.res.write(`data: ${JSON.stringify(data)}\n\n`);
			}
		}
	}

	// 修改完代码以后出现这样的现象
	// 正常情况下，AI 请求完成（不中止），前端接收到[DONE] ，会执行 es.close()，服务端可以监听到
	// 但如果中途中止请求，服务端就监听不到 req on close —— 奇了怪了！！！
	/* ctx.req.on('close', () => {
    console.log('req close...')
  }) */
});

export default router;
